#include "nav_utils/manual_traj_optimizer.h"

namespace Planner
{
    ManualTrajOptimizer::ManualTrajOptimizer(const ros::NodeHandle &node_handle)
        : nh_(node_handle)
    {
        init();
    }

    void ManualTrajOptimizer::init()
    {
        nh_.param("manual_traj_optimizer/optimize", optimize_flag_, false);
        nh_.param("manual_traj_optimizer/opt_seg_num", opt_seg_num_, 10);
        nh_.param("manual_traj_optimizer/manual_bagname", manual_bagname_, string(""));
        nh_.param("manual_traj_optimizer/optimize_bagname", optimize_bagname_, string(""));
        nh_.param("manual_traj_optimizer/cloud_filename", cloud_filename_, string(""));
        nh_.param("manual_traj_optimizer/init_height", init_height_, 1.4);
        nh_.param("manual_traj_optimizer/filter_num", filter_num_, 0);
        nh_.param("manual_traj_optimizer/filter_radius", filter_radius_, 0.1);
        nh_.param("manual_traj_optimizer/des_gap", des_gap_, 1.5);

        record_traj_pub_ = nh_.advertise<nav_msgs::Path>("record_traj_vis", 1);
        optimized_traj_pub_ = nh_.advertise<nav_msgs::Path>("optimized_traj_vis", 1);
        cloud_map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("cloud_map_vis", 1);
        cloud_map_filtered_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("cloud_map_filter_vis", 1);

        cloud_map_ = boost::make_shared<pcl::PointCloud<pcl::PointXYZ>>();
        cloud_map_filtered_ = boost::make_shared<pcl::PointCloud<pcl::PointXYZ>>();
        kdtree_ = boost::make_shared<pcl::KdTreeFLANN<pcl::PointXYZ>>();

        int bspline_order;
        nh_.param("manual_traj_optimizer/bspline_order", bspline_order, 3);
        bspline_smooth_.setParams(bspline_order, des_gap_);

        double speed;
        nh_.param("trajectory_optimization/order", order_, 5);
        nh_.param("trajectory_optimization/speed", speed, 1.0);
        traj_gen_.setParams(order_, speed);

        traj_opt_.setParams(nh_);
    }

    int ManualTrajOptimizer::optimize()
    {
        int ret = readData();
        if (ret < 0)
        {
            ROS_ERROR("input data invaild");
            return ret;
        }

        waypointsCheck();

        if (!optimize_flag_)
        {
            simpleSmooth();
        }
        else
        {
            multiFuncOptimize();
        }
        ROS_INFO("Trajectory Optimized");

        trajectoryCheck();

        writeData();
    }

    int ManualTrajOptimizer::readData()
    {
        vector<nav_msgs::Path::Ptr> traj_data;
        try
        {
            rosbag::Bag bag;
            bag.open(manual_bagname_, rosbag::bagmode::Read);
            rosbag::View view(bag, rosbag::TopicQuery("record_trag"));

            for (rosbag::MessageInstance const m : view)
            {
                nav_msgs::Path::Ptr s = m.instantiate<nav_msgs::Path>();
                if (s != nullptr)
                    traj_data.push_back(s);
            }
        }
        catch (const std::exception &e)
        {
            std::cerr << e.what() << '\n';
            return -1;
        }

        if (traj_data.size() != 1)
        {
            ROS_ERROR("Global trajectory msgs should obtain 1 msg, but now msg size is %d", traj_data.size());
            return 0;
        }

        if (pcl::io::loadPCDFile(cloud_filename_, *cloud_map_) == -1)
        {
            ROS_ERROR("Cannot read pcd file");
            return -1;
        }
        waypoints_.push_back(Vector3d(0, 0, init_height_));
        for (size_t i = 0; i < traj_data[0]->poses.size(); i++)
        {
            double x = traj_data[0]->poses[i].pose.position.x;
            double y = traj_data[0]->poses[i].pose.position.y;
            double z = traj_data[0]->poses[i].pose.position.z;
            if (waypoints_.size() == 1 && (x * x + y * y) < des_gap_)
                continue;
            waypoints_.push_back(Vector3d(x, y, z));
        }
        vector<Vector3d> smooth_temp;
        bspline_smooth_.smoothPath(waypoints_, smooth_temp);
        bspline_smooth_.samplePath(smooth_temp, waypoints_smoothed_);

        pcl::RadiusOutlierRemoval<pcl::PointXYZ> outrem;
        outrem.setInputCloud(cloud_map_);
        outrem.setRadiusSearch(filter_radius_);
        outrem.setMinNeighborsInRadius(filter_num_);
        outrem.filter(*cloud_map_filtered_);
        cout << "cloudMap size: " << cloud_map_->size() << ", after filtered: "
             << cloud_map_filtered_->size() << endl;

        kdtree_->setInputCloud(cloud_map_filtered_);

        pcl::toROSMsg(*cloud_map_, cloud_map_msgs_);
        cloud_map_msgs_.header.frame_id = "world";
        pcl::toROSMsg(*cloud_map_filtered_, cloud_map_filtered_msgs_);
        cloud_map_filtered_msgs_.header.frame_id = "world";
    }

    void ManualTrajOptimizer::writeData()
    {
        quad_msgs::Trajectory traj_msgs;
        coeffsToGlobalTrajMsg(coeffs_, time_seq_, traj_msgs);
        rosbag::Bag bag;
        bag.open(optimize_bagname_, rosbag::bagmode::Write);
        bag.write("global_trajectory", ros::Time::now(), traj_msgs);
        ROS_INFO("Write optimize trajectory to rosbag");
    }

    void ManualTrajOptimizer::simpleSmooth()
    {
        traj_gen_.setWayPoints(waypoints_smoothed_);
        traj_gen_.setBoundaryCondition(Vector3d(0, 0, 0), Vector3d(0, 0, 0), Vector3d(0, 0, 0), Vector3d(0, 0, 0));
        traj_gen_.solve(solve_mode::JERK);
        traj_gen_.getCoeffs(coeffs_);
        traj_gen_.getTimeSeq(time_seq_);
    }

    void ManualTrajOptimizer::multiFuncOptimize()
    {
        traj_opt_.setCloudMap(cloud_map_filtered_, kdtree_);

        int opt_times = waypoints_smoothed_.size() / opt_seg_num_;
        if (opt_times == 0)
            opt_times = 1;
        coeffs_ = vector<MatrixXd>(3, MatrixXd::Zero(waypoints_smoothed_.size() - 1, order_ + 1));
        time_seq_.clear();

        for (int s = 0; s < opt_times; s++)
        {
            int start_index = s * opt_seg_num_;
            int end_index = (s + 1) * opt_seg_num_;
            if (s == opt_times - 1)
                end_index = waypoints_smoothed_.size() - 1;
            cout << "optimizing trajectory from waypoints index " << start_index
                 << " to index " << end_index << endl;

            vector<Vector3d> input_waypoints(waypoints_smoothed_.begin() + start_index,
                                             waypoints_smoothed_.begin() + end_index + 1);
            vector<MatrixXd> part_coeffs;
            vector<double> part_time_seq;
            traj_opt_.setWayPoints(input_waypoints);
            Vector3d startVel(0, 0, 0), startAcc(0, 0, 0), goalVel(0, 0, 0), goalAcc(0, 0, 0);
            // 非起点startVel速度绝对值为1
            if (s > 0)
                startVel = (waypoints_smoothed_[start_index + 1] - waypoints_smoothed_[start_index]).normalized();
            // 非终点处速度绝对值为1
            if (s < opt_times - 1)
                goalVel = (waypoints_smoothed_[end_index + 1] - waypoints_smoothed_[end_index]).normalized();

            traj_opt_.setBoundaryCondition(startVel, startAcc, goalVel, goalAcc);
            if (traj_opt_.optimize() == -1) // -1为nlopt::FAILURE
            {
                ROS_ERROR("Optimizing trajectory failed");
                return;
            }

            traj_opt_.getCoeffs(part_coeffs);
            traj_opt_.getTimeSeq(part_time_seq);

            // 输出优化结果
            vector<double> cost_val;
            traj_opt_.getCostVal(cost_val);
            cout << "cost_time: " << cost_val[0] << ", cost_smooth: " << cost_val[1] << ", cost_collision: " << cost_val[2]
                 << ", cost_dynamic_vel: " << cost_val[3] << ", cost_dynamic_acc: " << cost_val[4] << endl;

            for (int i = 0; i < 3; i++)
                coeffs_[i].block(start_index, 0, end_index - start_index, 6) = part_coeffs[i];
            time_seq_.insert(time_seq_.end(), part_time_seq.begin(), part_time_seq.end());
        }
    }

    bool ManualTrajOptimizer::waypointsCheck()
    {
        bool flag = true;
        for (auto w : waypoints_)
        {
            double obsDis = calcDis(w[0], w[1], w[2]);
            if (obsDis < 1.4)
            {
                cout << "collision at (" << w[0] << ", " << w[1] << ", " << w[2] << "), dis: " << obsDis << endl;
                flag = false;
            }
        }
        if (flag)
            ROS_INFO("Waypoints Safety Check Passed");
        else
            ROS_WARN("Waypoints Safety Check failed");
        return flag;
    }

    bool ManualTrajOptimizer::trajectoryCheck()
    {
        if (time_seq_.empty())
            return false;
        int segments = time_seq_.size();
        int order = coeffs_[0].cols();
        bool flag = true;
        for (size_t s = 0; s < segments; s++)
        {
            double duration = time_seq_[s];
            for (double t = 0; t < duration; t += 0.2)
            {
                double x = 0, y = 0, z = 0;
                for (int r = 0; r < order; r++)
                {
                    x += coeffs_[0](s, r) * pow(t, r);
                    y += coeffs_[1](s, r) * pow(t, r);
                    z += coeffs_[2](s, r) * pow(t, r);
                }

                double obsDis = calcDis(x, y, z);
                if (obsDis < 1.4)
                {
                    cout << "collision at (" << x << ", " << y << ", " << z << "), dis: " << obsDis << endl;
                    flag = false;
                }
            }
        }
        if (flag)
            ROS_INFO("Trajectory Safety Check Passed");
        else
            ROS_WARN("Trajectory Safety Check failed");
        return flag;
    }

    double ManualTrajOptimizer::calcDis(double x, double y, double z)
    {
        pcl::PointXYZ pt(x, y, z);
        std::vector<int> pointIdxNKNSearch(1);
        std::vector<float> pointNKNSquaredDistance(1);
        if (kdtree_->nearestKSearch(pt, 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0)
        {
            float squaredDis = pointNKNSquaredDistance[0];
            double obsDis = std::sqrt(squaredDis);
            return obsDis;
        }
        std::cout << "kdtree search error\n";
        return -1;
    }

    void ManualTrajOptimizer::display()
    {
        cloud_map_pub_.publish(cloud_map_msgs_);
        cloud_map_filtered_pub_.publish(cloud_map_filtered_msgs_);
        displayTrajectory(waypoints_, record_traj_pub_);
        displayTrajectory(coeffs_, time_seq_, optimized_traj_pub_);
    }

} // namespace Planner
